#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import pandas as pd
from openpyxl.styles import Alignment, Font


# 定义颜色代码
class Colors:
    GREEN = '\033[92m'  # 绿色
    RED = '\033[91m'  # 红色
    YELLOW = '\033[93m'  # 黄色
    ENDC = '\033[0m'  # 结束颜色


def calculate_error(predicted, actual):
    """计算误差值"""
    return abs((predicted - actual) / actual * 100)


def is_accurate(predicted, actual):
    """判断预测是否准确"""
    if actual <= 5:
        return calculate_error(predicted, actual) <= 5
    elif actual <= 10:
        return calculate_error(predicted, actual) <= 10
    return False


def process_results(filename):
    """处理结果文件"""
    # 用于统计的数据
    total_5m = 0  # 5米及以下的样本总数
    accurate_5m = 0  # 5米及以下预测准确的数量
    total_5_10m = 0  # 5-10米的样本总数
    accurate_5_10m = 0  # 5-10米预测准确的数量

    # 定义列宽
    col_widths = {
        'filename': 15,
        'predicted': 10,
        'actual': 10,
        'error': 10
    }

    # 创建表头格式
    header_format = "{:<{filename_width}} {:<{predicted_width}} {:<{actual_width}} {:<{error_width}}"
    # 创建数据行格式 - 修改这里，不再限制小数位数
    data_format = "{:<{filename_width}} {:<{predicted_width}} {:<{actual_width}} {:<{error_width}.2f}"

    # 打印表头
    print(header_format.format(
        "文件名", "预测距离", "实际距离", "误差值(%)",
        filename_width=col_widths['filename'],
        predicted_width=col_widths['predicted'],
        actual_width=col_widths['actual'],
        error_width=col_widths['error']
    ))
    print("-" * (sum(col_widths.values()) + 3))  # +3 是因为有3个空格分隔符

    # 用于存储Excel数据
    excel_data = []

    with open(filename, 'r') as f:
        for line in f:
            # 跳过空行
            if not line.strip():
                continue

            # 解析数据
            parts = line.strip().split('\t')
            if len(parts) != 3:
                continue

            filename_item, pred_dist, actual_dist = parts
            pred_dist = float(pred_dist)
            actual_dist = float(actual_dist)

            # 计算误差值
            error = calculate_error(pred_dist, actual_dist)

            # 判断是否准确并选择颜色
            is_correct = is_accurate(pred_dist, actual_dist)
            if actual_dist > 10:
                color = Colors.YELLOW  # 大于10米用黄色
            else:
                color = Colors.GREEN if is_correct else Colors.RED

            # 输出每行的误差值（带颜色）
            formatted_line = data_format.format(
                filename_item, str(pred_dist), str(actual_dist), error,
                filename_width=col_widths['filename'],
                predicted_width=col_widths['predicted'],
                actual_width=col_widths['actual'],
                error_width=col_widths['error']
            )
            print(f"{color}{formatted_line}{Colors.ENDC}")

            # 将数据添加到Excel数据列表
            excel_data.append({
                "文件名": filename_item,
                "预测距离": pred_dist,
                "实际距离": actual_dist,
                "误差值(%)": round(error, 2)
            })

            # 统计准确率
            if actual_dist <= 5:
                total_5m += 1
                if is_correct:
                    accurate_5m += 1
            elif actual_dist <= 10:
                total_5_10m += 1
                if is_correct:
                    accurate_5_10m += 1

    # 计算准确率
    accuracy_5m = (accurate_5m / total_5m) * 100 if total_5m > 0 else 0
    accuracy_5_10m = (accurate_5_10m / total_5_10m) * 100 if total_5_10m > 0 else 0

    # 输出准确率
    print("\n准确率统计:")
    print("-" * (sum(col_widths.values()) + 3))
    if total_5m > 0:
        print(f"实际距离 ≤ 5m 的预测准确率: {accuracy_5m:.2f}%")

    if total_5_10m > 0:
        print(f"5m < 实际距离 ≤ 10m 的预测准确率: {accuracy_5_10m:.2f}%")

    # 创建Excel文件
    # 获取不带扩展名的输入文件名
    base_filename = os.path.splitext(filename)[0]
    excel_filename = f"{base_filename}.xlsx"

    # 创建DataFrame
    df = pd.DataFrame(excel_data)

    # 创建Excel写入器
    with pd.ExcelWriter(excel_filename, engine='openpyxl') as writer:
        # 写入数据，从第5行开始
        df.to_excel(writer, sheet_name='结果', startrow=4, index=False)

        # 获取工作表
        worksheet = writer.sheets['结果']

        # 设置单元格样式函数
        def set_cell_style(cell, value, is_header=False, center_align=False):
            cell.value = value
            if is_header:
                cell.font = Font(bold=True)
            if center_align:
                cell.alignment = Alignment(horizontal='center')

        # 写入准确率统计结果（占用四列空间）
        # 第一行 - 准确率统计标题
        worksheet.merge_cells('A1:D1')
        header_cell = worksheet.cell(row=1, column=1)
        set_cell_style(header_cell, "准确率统计", is_header=True, center_align=True)

        # 准确率行的格式统一
        # 第二行 - 5米以内准确率
        worksheet.merge_cells('A2:B2')
        label_cell_1 = worksheet.cell(row=2, column=1)
        set_cell_style(label_cell_1, "实际距离 ≤ 5m 的预测准确率")

        worksheet.merge_cells('C2:D2')
        value_cell_1 = worksheet.cell(row=2, column=3)
        set_cell_style(value_cell_1, f"{accuracy_5m:.2f}%", center_align=True)

        # 第三行 - 5-10米准确率
        worksheet.merge_cells('A3:B3')
        label_cell_2 = worksheet.cell(row=3, column=1)
        set_cell_style(label_cell_2, "5m < 实际距离 ≤ 10m 的预测准确率")

        worksheet.merge_cells('C3:D3')
        value_cell_2 = worksheet.cell(row=3, column=3)
        set_cell_style(value_cell_2, f"{accuracy_5_10m:.2f}%", center_align=True)

        # 空行
        worksheet.merge_cells('A4:D4')

    print(f"\nExcel结果已保存至: {excel_filename}")


if __name__ == "__main__":
    process_results("resultAll.txt")